
import os, glob
import torch
from torch import nn
from model import makeResNet152, makeResNet50
from torch import optim
from torch.utils.data import DataLoader
from dataset import classiferDataset, VeRiClassiferDataset
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms
import csv

def getFeatureMap(modelPath, imgsListGlobPath, featureMapPath):

    # model = makeResNet152(576)
    # model = torch.load(".\\savedModels\\model_epoch_70_right_0.999602943512097.pkl")
    model = torch.load(modelPath)
    model = nn.Sequential(
        *list(model.children())[:-1],
    )
    print(model)

    tf = transforms.Compose([
        lambda x: Image.open(x).convert('RGB'),
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406],
        #                     std=[0.229, 0.224, 0.225])
    ])

    # imgsList = glob.glob(".\\data\\VeRi\\image_test\\*.jpg")
    # featureMapPath = ".\\data\\VeRi\\image_test_featuremap\\"
    # imgsList = glob.glob(".\\data\\VeRi\\image_query\\*.jpg")
    # featureMapPath = ".\\data\\VeRi\\image_query_featuremap\\"
    imgsList = glob.glob(imgsListGlobPath)

    for img in imgsList:

        x = tf(img)
        x = x.unsqueeze(0).to(torch.device("cuda"))
        y = model(x)
        y = y.squeeze(0)
        # print(img)
        # print(".\\data\\VeRi\\image_test_featuremap\\" + img.split("\\")[-1][:-3] + "pt")
        torch.save(y, featureMapPath + img.split("\\")[-1][:-3] + "pt")



def compareFeaturemap():
    queryPath = ".\\data\\VeRi\\image_query_featuremap\\"
    testPath = ".\\data\\VeRi\\image_test_featuremap\\"

    queryPathList = glob.glob(queryPath + "*.pt")
    testPathList = glob.glob(testPath + "*.pt")

    csvPath = ".\\data\\VeRi\\query_csv\\"

    for query in queryPathList:
        f = open(csvPath + query.split("\\")[-1][:-2] + "csv", "w", newline="")
        writer = csv.writer(f)

        queryTensor = torch.load(query)

        res = []
        for test in testPathList:
            testTensor = torch.load(test)
            distance = (queryTensor - testTensor).pow(2).sum(0)
            distance = float(distance)

            # writer.writerow((test.split("\\")[-1], distance))
            res.append((test.split("\\")[-1], distance))

        res = sorted(res, key=lambda x: x[1])
        for row in res:
            writer.writerow((row[0], row[1]))

        f.close()



def csvToChart():
    queryImgPath = ".\\data\\VeRi\\image_query\\"
    testImgPath = ".\\data\\VeRi\\image_test\\"
    # csvPath = ".\\data\\VeRi\\query_csv\\"
    csvPath = ".\\data\\VeRi\\result\\query_csv\\"

    csvFileNameList = glob.glob(csvPath + "*.csv")


    queryNum = 0
    queryRight = 0

    for csvFile in csvFileNameList:
        # print("======================================================================")
        print(csvFile)
        queryImgName = queryImgPath + csvFile.split("\\")[-1].split(".")[0] + ".jpg"
        queryImgClass = csvFile.split("\\")[-1][:4]

        f = open(csvFile, "r")
        reader = csv.reader(f)

        distanceList = []
        for row in reader:
            distanceList.append((row[0], float(row[1])))

        distanceList = sorted(distanceList, key=lambda x: x[1])

        totalNum = 1
        totalRight = 0

        for i in range(1, totalNum + 1):
            testImgName = distanceList[i][0]

            testImgClass = testImgName[:4]
            testImgPath = testImgPath + testImgName.split(".")[0] + ".jpg"
            # print(testImgClass)
            # print(testImgName)
            if testImgClass == queryImgClass:
                totalRight += 1

        queryRight += totalRight
        queryNum += totalNum

        if totalNum != totalRight:
            print(totalRight / totalNum)

        # print("======================================================================")

    print(queryRight / queryNum)



def csvToChart2():
    queryImgPath = ".\\data\\VeRi\\image_query\\"
    testImgPath = ".\\data\\VeRi\\image_test\\"
    csvPath = ".\\data\\VeRi\\query_csv\\"

    csvFileNameList = glob.glob(csvPath + "*.csv")

    # get every class of image number
    everyClassNumber = {}
    allTestImgs = glob.glob(testImgPath + "*.jpg")
    for imgName in allTestImgs:
        imgClass = imgName.split("\\")[-1][:4]
        if imgClass in everyClassNumber.keys():
            everyClassNumber[imgClass] += 1
        else:
            # everyClassNumber[imgClass] = 1
            everyClassNumber[imgClass] = 0      # remove query itself

    # start get all mAP
    mAPSum = 0
    mAPNum = 0

    for csvFile in csvFileNameList:
        # print("======================================================================")
        # print(csvFile)
        queryImgName = queryImgPath + csvFile.split("\\")[-1].split(".")[0] + ".jpg"
        queryImgClass = csvFile.split("\\")[-1][:4]
        queryClassNumber = everyClassNumber[queryImgClass]

        f = open(csvFile, "r")
        reader = csv.reader(f)


        # get distance between queryImg and all testImgs, and sort
        distanceList = []
        for row in reader:
            distanceList.append((row[0], float(row[1])))
        distanceList = sorted(distanceList, key=lambda x: x[1])


        # traversing the sorted distanceList
        totalNum = 0
        totalAcc = 0
        accuracyList = []    # n index: Accuracy of the first n queried results
        accIndexList = []    # n index: distanceList[n] is right or false

        for i in range(0, len(distanceList)):
            testImgName = distanceList[i][0]

            testImgClass = testImgName[:4]
            testImgPath = testImgPath + testImgName.split(".")[0] + ".jpg"
            # print(testImgClass)
            # print(testImgName)

            if testImgClass == queryImgClass:
                totalAcc += 1
                accIndexList.append(True)
            else:
                accIndexList.append(False)
            totalNum += 1

            accuracyList.append(totalAcc/totalNum)

            # print(i,": ", accuracyList[-1], "\t\t\t\t", accIndexList[-1])

            if totalAcc == queryClassNumber:
                break

        # calculate the query mAP
        # except for the first element, because it's query itself
        queryClassNumber -= 1
        accuracyList[0] = accuracyList[1]

        mAP = 0
        for i in range(1, len(accuracyList)):
            if not accIndexList[i]:
                continue

            mAP += ((accuracyList[i] + accuracyList[i - 1]) / 2 / queryClassNumber)

        mAPSum += mAP
        mAPNum += 1
        print("this mAP: ", mAP, "   total mAP: ", mAPSum / mAPNum)

        # print("======================================================================")

        # print(queryClassNumber)
        # break


    print("final mAP: ", mAPSum / mAPNum)




def csvToChart3():
    queryImgPath = ".\\data\\VeRi\\image_query\\"
    testImgPath = ".\\data\\VeRi\\image_test\\"
    csvPath = ".\\data\\VeRi\\query_csv\\"
    savePltPath = ".\\data\\VeRi\\query_jpg\\"

    csvFileNameList = glob.glob(csvPath + "*.csv")

    # get every class of image number
    everyClassNumber = {}
    allTestImgs = glob.glob(testImgPath + "*.jpg")
    for imgName in allTestImgs:
        imgClass = imgName.split("\\")[-1][:4]
        if imgClass in everyClassNumber.keys():
            everyClassNumber[imgClass] += 1
        else:
            # everyClassNumber[imgClass] = 1
            everyClassNumber[imgClass] = 0      # remove query itself

    # start get all mAP
    mAPSum = 0
    mAPNum = 0

    for csvFile in csvFileNameList:
        # print("======================================================================")
        # print(csvFile)
        queryImgName = queryImgPath + csvFile.split("\\")[-1].split(".")[0] + ".jpg"
        queryPltResSaveName = savePltPath + csvFile.split("\\")[-1].split(".")[0] + ".jpg"
        queryImgClass = csvFile.split("\\")[-1][:4]
        queryClassNumber = everyClassNumber[queryImgClass]

        f = open(csvFile, "r")
        reader = csv.reader(f)


        # get distance between queryImg and all testImgs, and sort
        distanceList = []
        for row in reader:
            distanceList.append((row[0], float(row[1])))
        distanceList = sorted(distanceList, key=lambda x: x[1])


        # traversing the sorted distanceList
        totalNum = 0
        totalAcc = 0
        accuracyList = []    # n index: Accuracy of the first n queried results
        accIndexList = []    # n index: distanceList[n] is right or false


        # plt init
        plt.figure()
        pltCol = 15
        pltRow = 7


        for i in range(0, len(distanceList)):
            testImgName = distanceList[i][0]

            testImgClass = testImgName[:4]
            testImgFullName = testImgPath + testImgName.split(".")[0] + ".jpg"
            # print(testImgClass)
            # print(testImgName)

            if testImgClass == queryImgClass:
                totalAcc += 1
                accIndexList.append(True)
            else:
                accIndexList.append(False)
            totalNum += 1

            accuracyList.append(totalAcc/totalNum)

            # print(i,": ", accuracyList[-1], "\t\t\t\t", accIndexList[-1])

            # make plt
            if i < 15*7:
                plt.subplot(pltRow, pltCol, i + 1)

                testImg = Image.open(testImgFullName)
                imgTitle = testImgFullName.split("\\")[-1][:9]
                if i == 0:
                    titleColor = "purple"
                    imgTitle = "query\n" + imgTitle
                else:
                    if testImgClass == queryImgClass:
                        titleColor = "green"
                    else:
                        titleColor = "red"
                    imgTitle = imgTitle + "\n" + str(distanceList[i][1])[:9]
                plt.title(imgTitle, color=titleColor)
                plt.imshow(testImg)
                plt.axis('off')

            if totalAcc == queryClassNumber:
                break


        # calculate the query mAP
        # except for the first element, because it's query itself
        queryClassNumber -= 1
        accuracyList[0] = accuracyList[1]

        mAP = 0
        for i in range(1, len(accuracyList)):
            if not accIndexList[i]:
                continue

            mAP += ((accuracyList[i] + accuracyList[i - 1]) / 2 / queryClassNumber)


        # plt write mAP and save
        plt.subplot(pltRow, pltCol, 1)
        plt.text(0, -150, "mAP: %.3f" % mAP, fontsize=30)
        plt.gcf().set_size_inches(2560 / 100, 1440 / 100)
        plt.savefig(queryPltResSaveName, dpi=300)

        mAPSum += mAP
        mAPNum += 1
        print("this mAP: ", mAP, "   total mAP: ", mAPSum / mAPNum)

        # print("======================================================================")

        # print(queryClassNumber)
        # break


    print("final mAP: ", mAPSum / mAPNum)




if __name__ == '__main__':

    modelPath = ".\\savedModels\\model_epoch_73_right_0.9996823548096776.pkl"


    model = torch.load(modelPath)
    model = nn.Sequential(
        *list(model.children())[:-1],
    )
    print(model)
    torch.save(model, "model.pt")

    # getFeatureMap(
    #     modelPath,
    #     ".\\data\\VeRi\\image_test\\*.jpg",
    #     ".\\data\\VeRi\\image_test_featuremap\\"
    # )
    # getFeatureMap(
    #     modelPath,
    #     ".\\data\\VeRi\\image_query\\*.jpg",
    #     ".\\data\\VeRi\\image_query_featuremap\\"
    # )
    # compareFeaturemap()


    # csvToChart()
    # csvToChart2()
    # csvToChart3()
    #665.31